package main

import (
	"flag"
	"github.com/miekg/dns"
	"github.com/rocketmq-exporter-go/config"
	"github.com/rocketmq-exporter-go/exporter"
	"net"
	"strconv"
	"strings"

	"github.com/apache/rocketmq-client-go/v2/rlog"
	"github.com/prometheus/common/version"
	"gopkg.in/alecthomas/kingpin.v2"
	// _ "net/http/pprof"
)

func toFlagInt(name string, help string, value string) *int {
	flag.CommandLine.String(name, value, help) // hack around flag.Parse and glog.init flags
	return kingpin.Flag(name, help).Default(value).Int()
}

func toFlagString(name string, help string, value string) *string {
	flag.CommandLine.String(name, value, help) // hack around flag.Parse and glog.init flags
	return kingpin.Flag(name, help).Default(value).String()
}

func toFlagStrings(name string, help string, value string) *[]string {
	flag.CommandLine.String(name, value, help) // hack around flag.Parse and glog.init flags
	return kingpin.Flag(name, help).Default(value).Strings()
}

func toFlagBool(name string, help string, value bool, valueString string) *bool {
	flag.CommandLine.Bool(name, value, help) // hack around flag.Parse and glog.init flags
	return kingpin.Flag(name, help).Default(valueString).Bool()
}

func toFlagIntVar(name string, help string, value int, valueString string, target *int) {
	flag.CommandLine.Int(name, value, help) // hack around flag.Parse and glog.init flags
	kingpin.Flag(name, help).Default(valueString).IntVar(target)
}

func main() {

	// 读取配置文件config.yaml，并解析为结构体
	conf := config.GetConfig()
	if conf == nil {

		return
	}

	var (
		listenAddress = &conf.Web.ListenAddress
		metricsPath   = &conf.Web.MetricsPath

		logFileEnable = &conf.Log.File.Enable
		logFilePath   = &conf.Log.File.Path

		useDomain               = &conf.RocketMQ.NameServer.UseDomain
		nameServerIpAddress     = &conf.RocketMQ.NameServer.NameServerIpAddress
		nameServerDomainAddress = &conf.RocketMQ.NameServer.NameServerDomain
		nameServerPort          = &conf.RocketMQ.NameServer.NameServerPort

		opts = &exporter.RocketmqOptions{}
	)

	listenAddress = toFlagString("web.listen-address", "Address to listen on for web interface and telemetry.", *listenAddress)
	metricsPath = toFlagString("web.telemetry-path", "Path under which to expose metrics.", *metricsPath)

	logFileEnable = toFlagBool("log.file.enable", "Log write to file enable.", *logFileEnable, strconv.FormatBool(*logFileEnable))
	logFilePath = toFlagString("log.file.path", "Path of log file.", *logFilePath)

	useDomain = toFlagBool("rocketmq.nameserver.use.domian", "Address of rocketmq nameserver whether use domian.", *useDomain, strconv.FormatBool(*useDomain))
	nsia := strings.Join(*nameServerIpAddress, ",")
	nameServerIpAddress = toFlagStrings("rocketmq.nameserver", "Address (host:port) of rocketmq nameserver.", nsia)
	nameServerDomainAddress = toFlagString("rocketmq.nameserver.domain", "Address (domain:port) of rocketmq nameserver.", *nameServerDomainAddress)
	nameServerPort = toFlagInt("rocketmq.nameserver.port", "Port of rocketmq nameserver.", strconv.Itoa(*nameServerPort))
	// 打印配置
	rlog.Info("Config ", map[string]interface{}{
		"listenAddress":           *listenAddress,
		"metricsPath":             *metricsPath,
		"logFileEnable":           *logFileEnable,
		"logFilePath":             *logFilePath,
		"useDomain":               *useDomain,
		"nameServerIpAddress":     *nameServerIpAddress,
		"nameServerDomainAddress": *nameServerDomainAddress,
		"nameServerPort":          *nameServerPort,
	})

	toFlagIntVar("workers", "Number of workers", 100, "100", &opts.Workers)

	kingpin.Version(version.Print("rocketmq_exporter"))
	kingpin.HelpFlag.Short('h')
	kingpin.Parse()

	if *logFileEnable {
		rlog.SetOutputPath(*logFilePath)
	}

	if *useDomain {

		var ips, err = LookupIP(*nameServerDomainAddress, *nameServerPort)
		rlog.Fatal("rocketmq.nameserver.domain parse err", map[string]interface{}{
			"domain": *nameServerDomainAddress,
			"err":    err,
		})
		opts.NameServerAddress = ips

	} else {
		opts.NameServerAddress = *nameServerIpAddress
	}

	rlog.Info("Fly ", map[string]interface{}{
		"listenAddress": *listenAddress,
		"metricsPath":   *metricsPath,
		"opts":          *opts,
	})

	exporter.Fly(*listenAddress, *metricsPath, opts)

}

// Because rocketmq client requires ip address rather than domain address
// LookupIP looks up host
// It returns a slice of that host's IPv4 addresses.
func LookupIP(host string, port int) ([]string, error) {
	ips, err := net.LookupIP(host)

	if err != nil {
		// 尝试第二种方式获取ip，适用于k8s环境
		// 设置 DNS 服务器的地址
		dnsServer := "192.168.100.3:53"
		// 执行 DNS 查询
		ipss, err := queryDNS(dnsServer, host)
		if err != nil {
			return nil, err
		}
		var ipList = make([]string, len(ips))
		for index, ip := range ipss {
			ipList[index] = ip + ":" + strconv.Itoa(port)
		}
		return ipList, nil
	}
	var ipList = make([]string, len(ips))
	for index, ip := range ips {
		ipList[index] = ip.To4().String() + ":" + strconv.Itoa(port)
	}

	return ipList, nil
}

// queryDNS 使用 miekg/dns 库进行 DNS 查询
func queryDNS(server, domain string) ([]string, error) {
	client := dns.Client{}
	msg := dns.Msg{}
	msg.SetQuestion(dns.Fqdn(domain), dns.TypeA)

	r, _, err := client.Exchange(&msg, server)
	if err != nil {
		return nil, err
	}

	var ips []string
	for _, ans := range r.Answer {
		if a, ok := ans.(*dns.A); ok {
			ips = append(ips, a.A.String())
		}
	}

	return ips, nil
}
